import torch.nn as nn
import torch


class MAEWithPrompt(nn.Module):
    """
    在原始 Masked Autoencoder (MAE) 模型中添加可学习的提示（Prompt）嵌入。
    通过将一组可学习的提示嵌入向量与输入数据拼接在一起，增强模型的表达能力。

    参数：
        original_model (nn.Module): 传入的原始模型（如 MAE 或其他模型），其将使用添加提示后的输入进行训练。
        prompt_dim (int): 提示嵌入的维度，决定了每个提示向量的大小。
        num_prompts (int): 提示的数量，决定了在输入中添加多少个提示嵌入。

    """

    def __init__(self, original_model, prompt_dim, num_prompts):
        # 初始化 MAEWithPrompt 类
        super(MAEWithPrompt, self).__init__()

        # 存储原始模型
        self.original_model = original_model

        # 初始化提示嵌入参数（可训练的）
        # 这里我们为每个提示初始化一个随机的向量，形状为 (num_prompts, prompt_dim)
        self.prompt_embeddings = nn.Parameter(torch.randn(num_prompts, prompt_dim))

    def forward(self, x):
        """
        前向传播函数，将提示嵌入与输入数据拼接，并传递给原始模型。

        参数：
            x (Tensor): 输入张量，通常是图像或特征，形状为 (batch_size, seq_length, feature_dim)。

        返回：
            Tensor: 经过拼接并输入原始模型后的输出。
        """
        # 获取当前批次的大小
        batch_size = x.shape[0]

        # 扩展提示嵌入，以匹配批次大小
        # `self.prompt_embeddings` 形状为 (num_prompts, prompt_dim)，
        # 通过 `unsqueeze(0)` 扩展为 (1, num_prompts, prompt_dim)，然后使用 `expand` 扩展为 (batch_size, num_prompts, prompt_dim)
        prompt = self.prompt_embeddings.unsqueeze(0).expand(batch_size, -1, -1)

        # 将提示嵌入与输入数据拼接
        # 假设输入 `x` 的形状为 (batch_size, seq_length, feature_dim)，
        # 通过 `torch.cat` 将提示嵌入与输入张量拼接，拼接维度为 1 (即 seq_length 维度)，
        # 所以拼接后的 `x_with_prompt` 形状为 (batch_size, seq_length + num_prompts, feature_dim)
        x_with_prompt = torch.cat([prompt, x], dim=1)

        # 将拼接后的数据传递给原始模型
        return self.original_model(x_with_prompt)

